from dataclasses import dataclass, field
from typing import Dict, List, Tuple
from .common import Simulator, CostCalculator
from abc import ABC, abstractmethod

class VocabParallelEmbeddingSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.seq_len = mp['input_shape'][0][0]
        self.batch_size = mp['input_shape'][0][1]
        self.hidden_size = mp['weight_shape'][1]
        self.weight_hidden_size = mp['weight_shape'][0]

    def get_calculate_workload(self, tp_size: int) -> List[float]:
        return [0,0,0,0,0]
    
    def get_communication_workload(self, tp_size):
        return 0

    def get_dp_communication_workload(self, tp_size, dp_size):
        if dp_size == 1:
            return 0
        # *2 参数和梯度, *4 4字节
        return self.weight_hidden_size * self.hidden_size * 2 * 4 / tp_size / 1024 / 1024 / 1024

class EmbeddingSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.weight_hidden_size = mp['weight_shape'][0]
        self.hidden_size = mp['weight_shape'][1]
    
    def get_calculate_workload(self, tp_size: int) -> List[float]:
        return [0,0,0,0,0]
    
    def get_communication_workload(self, tp_size):
        return 0
    
    # *2参数和梯度 *4 4字节
    def get_dp_communication_workload(self, tp_size, dp_size):
        if dp_size == 1:
            return 0
        return self.weight_hidden_size * self.hidden_size * 2 * 4 / tp_size / 1024 / 1024 / 1024

class EmbeddingDropoutSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.seq_len = mp['input_shape'][0][0]
        self.batch_size = mp['input_shape'][0][1]
        self.hidden_size = mp['input_shape'][0][2]
    
    def get_calculate_workload(self, tp_size: int) -> List[float]:
        return [0,0,0,0,0]
    
    # embedding 层 all-reduce 通信, 一次2fi，前向反响各一次，共4 fi, 4字节
    def get_communication_workload(self, tp_size):
        if tp_size == 1: return 0

        return self.seq_len * self.batch_size * self.hidden_size * 4 * 4 / 1024 / 1024 / 1024
    
    # *2 参数和梯度， 4字节
    def get_dp_communication_workload(self, tp_size, dp_size):
        if dp_size == 1: return 0
        return self.weight_hidden_size * self.hidden_size / tp_size * 2 * 4 / 1024 / 1024 / 1024
    
class RotaryEmbeddingSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
    
    # 计算量下移到Attetnion，因为input_shape没记录下来
    def get_calculate_workload(self, tp_size: int) -> List[float]:
        return [0,0,0,0,0]
    
    def get_communication_workload(self, tp_size):
        return 0
    
    def get_dp_communication_workload(self, tp_size, dp_size):
        return 0


@dataclass
class LanguageModelEmbedding(CostCalculator):
    word_embeddings: VocabParallelEmbeddingSimulator = None
    position_embeddings: EmbeddingSimulator = None
    tokentype_embeddings: EmbeddingSimulator = None
    embedding_dropout: EmbeddingDropoutSimulator = None

    def statistic_single_layer_memory_workload(self, tp_size) -> Tuple:
        theoretical_embedding_memory = self.word_embeddings.weight_hidden_size * self.word_embeddings.hidden_size * 4 / 1024 / 1024 / 1024 / tp_size

        embedding_active_memory = (self.embedding_dropout.mp['max_memory_allocated'] - self.embedding_dropout.mp['init_memory_allocated']) / 1024 / 1024 /1024 / tp_size
        return theoretical_embedding_memory, embedding_active_memory

class GPTLanguageModelEmbeddingAbstractBuilder(ABC):
    @abstractmethod
    def build_word_embeddings(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_position_embeddings(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_tokentype_embeddings(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_embedding_dropout(self, args: Dict[str, object]): ...

class LlamaLanguageModelEmbeddingAbstractBuilder(ABC):
    @abstractmethod
    def build_rotary_embedding(self, args: Dict[str, object]): ...

class LanguageModelEmbeddingAbstractBuilder(GPTLanguageModelEmbeddingAbstractBuilder, LlamaLanguageModelEmbeddingAbstractBuilder):
    @abstractmethod
    def build_language_embeddings(self,args: Dict[str, object]): ...